feat: Add NETopKV function. #1251

morgolock · 2026-01-15T17:38:11Z

The Neon(TM) implementation of TopKV reduces execution time from 447.8 ms (scalar CPP) to 11.65 ms for the same workload (F32, C=1000, N=32000, k=3, 6 threads), achieving an approximate 38× speedup. This gain comes from SIMD vectorization, removal of per-element branches, and a more efficient inner loop.
Resolves ARMCL-1227

Change-Id: Ifdf161ce4254dc5ecd57aff9ae22410facd31705

arm_compute/runtime/NEON/functions/NETopKV.h

gunes-arm · 2026-01-15T20:34:41Z

arm_compute/runtime/NEON/functions/NETopKV.h

+     */
+    void configure(const ITensor *predictions, const ITensor *targets, ITensor *output, const unsigned int k);
+
+    /** Static function to check if given info will lead to a valid configuration of @ref CPPTopKVKernel


CpuTopKVKernel

Done in next patchset

src/cpu/kernels/CpuTopKVKernel.cpp

gunes-arm · 2026-01-15T20:48:43Z

src/cpu/operators/CpuTopKV.h

+    void
+    configure(const ITensorInfo *predictions, const ITensorInfo *targets, ITensorInfo *output, const unsigned int k);
+
+    /** Static function to check if given info will lead to a valid configuration of @ref CPPTopKVKernel


can we just say "similar to @ref CpuTopKV::configure()" to reduce duplication? Have a look at arm_compute/runtime/experimental/operators/CpuActivation.h

Same comment is valid for every header

Done in next patchset

src/cpu/operators/CpuTopKV.cpp

tests/validation/reference/TopKV.cpp

tests/validation/fixtures/TopKVLayerFixture.h

gunes-arm · 2026-01-15T22:57:08Z

tests/validation/fixtures/TopKVLayerFixture.h

+#define ACL_TESTS_VALIDATION_FIXTURES_TOPKVLAYERFIXTURE_H
+
+#include "arm_compute/core/TensorShape.h"
+#include "arm_compute/core/Types.h"


Do we need anything particular from this file?

Removed them in next patch

gunes-arm · 2026-01-18T13:39:20Z

src/cpu/kernels/CpuTopKVKernel.h

+
+#include "src/core/common/Macros.h"
+#include "src/cpu/ICpuKernel.h"
+


Looking at more carefully, I think we need to include the STL vector, string and cstdint headers in this file; not inside the corresponding cpp. Also, type_traits for std::add_pointer.

Same goes for ITensorInfo.h, ITensor.h.

We also need to include

arm_compute/core/Window.h for Window

arm_compute/core/CPP/CPPTypes.h for CPUInfo

src/cpu/kernels/CpuKernelSelectionTypes.h for DataTypeSelector.. stuff

arm_compute/core/Error.h for Status

In a nutshell, this file should be self-sufficient and we don't need to repeat the same includes in the corresponding .cpp file for brevity.

Can you please review this file and CpuTopKVKernek.cpp accordingly?

Updated in next patch

string and type_traits are not added here.
We should also add ITensorInfo.h, ITensor.h here and remove from CpuTopKVKernel.cpp.

src/cpu/kernels/CpuTopKVKernel.h

tests/validation/reference/TopKV.cpp

src/cpu/kernels/topkv/generic/neon/fp16.cpp

src/cpu/operators/CpuTopKV.h

gunes-arm · 2026-01-18T15:32:23Z

tests/validation/fixtures/TopKVLayerFixture.h

+
+            for (unsigned int n = 0; n < N; ++n)
+            {
+                const uint32_t target = *reinterpret_cast<const uint32_t *>(targets(Coordinates{static_cast<int>(n)}));


targets[i]??

Does this not work?

tests/validation/reference/TopKV.cpp

tests/validation/reference/TopKV.h

arm_compute/runtime/NEON/functions/NETopKV.h

* The Neon(TM) implementation of TopKV reduces execution time from 447.8 ms (CPP) to 11.65 ms for the same workload (F32, C=1000, N=32000, k=3, 6 threads), achieving an approximate 38× speedup. This gain comes from SIMD vectorization, removal of per-element branches, and a more efficient inner loop * Resolves ARMCL-1227. Change-Id: Ifdf161ce4254dc5ecd57aff9ae22410facd31705 Signed-off-by: Pablo Marquez Tello <pablo.tello@arm.com>

gunes-arm · 2026-01-23T11:47:46Z

docs/user_guide/operator_list.dox

Copyright should be 2026

gunes-arm · 2026-01-23T11:52:12Z

src/cpu/kernels/CpuTopKVKernel.h

+
+#include "src/core/common/Macros.h"
+#include "src/cpu/ICpuKernel.h"
+


string and type_traits are not added here.
We should also add ITensorInfo.h, ITensor.h here and remove from CpuTopKVKernel.cpp.

gunes-arm · 2026-01-23T11:57:49Z

src/cpu/kernels/topkv/generic/neon/fp16.cpp

+    // Load 8 fp16
+    const float16x8_t v16 = vld1q_f16(reinterpret_cast<const __fp16 *>(ptr));
+
+    // Compare in fp32 (same correctness story you already debugged)


What do we mean by "you already debugged" here?

gunes-arm · 2026-01-23T12:02:28Z

src/cpu/kernels/topkv/generic/neon/fp32.cpp

+
+} // namespace detail
+
+void topkv_fp32_neon(const ITensor *in1, const ITensor *in2, ITensor *out, uint32_t k, const Window &win)


Can we change these to predictions & targets in all classes and functions? (Except the operator list.dox unfortunately due to the convention there). There is a mix-up of in1/src0 in different files.

gunes-arm · 2026-01-23T12:05:16Z

arm_compute/runtime/NEON/functions/NETopKV.h

+     * |S32            |U32            |U8             |
+     * |F16            |U32            |U8             |
+     * |F32            |U32            |U8             |
+     * |U8             |U32            |U8             |


I believe this U8 line extra here

gunes-arm · 2026-01-23T12:11:45Z

tests/datasets/TopKVDataset.h

+
+#include "tests/framework/datasets/Datasets.h"
+
+#include <type_traits>


Why is type_traits here? I think we only need , am I missing anything?

gunes-arm · 2026-01-23T12:13:30Z

tests/validation/NEON/TopKV.cpp

+
+TEST_SUITE(NEON)
+TEST_SUITE(TopKVLayer)
+// clang-format on


I don't think we need these two lines.

gunes-arm · 2026-01-23T12:16:52Z

tests/validation/NEON/TopKV.cpp

+TEST_SUITE(TopKVLayer)
+
+template <typename T>
+using CPPTopKVLayerFixture = TopKVValidationFixture<Tensor, Accessor, CPPTopKV, T>;


I see that it was useful when comparing this implementation but we never have a CPP test suite in the NEON/ directory. The build system might miss to compile this test because we might ignore (we may already be ignoring) any files with under tests/validation/NEON when neon=0 option is provided to Scons.

I think these tests are valuable. Let's keep them but put them under tests/validation/CPP/TopKV.cpp

gunes-arm · 2026-01-23T12:17:51Z

tests/validation/fixtures/TopKVLayerFixture.h

+
+            for (unsigned int n = 0; n < N; ++n)
+            {
+                const uint32_t target = *reinterpret_cast<const uint32_t *>(targets(Coordinates{static_cast<int>(n)}));


Does this not work?

gunes-arm · 2026-01-23T12:22:38Z

src/cpu/kernels/CpuTopKVKernel.h

+    /** Initialise the kernel's input, dst and border mode.
+     *
+     * @param[in]  src0 First input tensor info. Data types supported: QASYMM8/QASYMM8_SIGNED/F16/F32/S32
+     * @param[in]  src1 Second input tensor info. Data types supported: U32


I think we're missing the argument k here

gunes-arm · 2026-01-23T12:30:00Z

src/cpu/kernels/CpuTopKVKernel.cpp

+
+    // targets must match batch
+    // targets is expected to contain N elements (shape [N])
+    ARM_COMPUTE_RETURN_ERROR_ON_MSG(src1.num_dimensions() < 1, "targets must have at least 1 dimension");


I think we only need to copy what's already inside CPPTopKVKernel.cpp:

ARM_COMPUTE_RETURN_ERROR_ON(predictions->num_dimensions() > 2); ARM_COMPUTE_RETURN_ERROR_ON(targets->num_dimensions() > 1);

morgolock force-pushed the pr/cpu_topkv branch from cb2af0c to 1f0778c Compare January 16, 2026 09:28

gunes-arm requested changes Jan 16, 2026

View reviewed changes

morgolock force-pushed the pr/cpu_topkv branch 4 times, most recently from 89cd2b8 to cd41038 Compare January 16, 2026 16:25

gunes-arm requested changes Jan 19, 2026

View reviewed changes

morgolock mentioned this pull request Jan 19, 2026

NEON implementation for TopK/TopKV #1227

Closed

morgolock force-pushed the pr/cpu_topkv branch 3 times, most recently from 541f913 to e0090dc Compare January 20, 2026 17:43

morgolock force-pushed the pr/cpu_topkv branch from e0090dc to 11238b6 Compare January 20, 2026 19:06

gunes-arm requested changes Jan 23, 2026

View reviewed changes

gunes-arm reviewed Jan 23, 2026

View reviewed changes


		#include "src/core/common/Macros.h"
		#include "src/cpu/ICpuKernel.h"


		} // namespace detail

		void topkv_fp32_neon(const ITensor in1, const ITensor in2, ITensor *out, uint32_t k, const Window &win)


		#include "tests/framework/datasets/Datasets.h"

		#include <type_traits>

feat: Add NETopKV function. #1251

Are you sure you want to change the base?

feat: Add NETopKV function. #1251

Conversation

morgolock commented Jan 15, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gunes-arm Jan 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

gunes-arm Jan 18, 2026 •

edited

Loading